#include "col2_im_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include <vector>
#include <cstdint>
#include <iostream>
#include <algorithm>

namespace optiling {
    static inline uint32_t AlignUp(uint32_t a, uint32_t b) 
    {
        if (b == 0)
            return a;
        return (a + b - 1) / b * b;
    }
    
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t CORE_NUM = 8;
    const uint32_t MIN_TILE_SIZE = 16;
    const uint32_t BUFFER_NUM = 2;
    const uint32_t VEC_LEN = 8;

    static ge::graphStatus TilingFunc(gert::TilingContext* context) {
        if (!context) {
            std::cout << "ERROR: TilingContext is null!" << std::endl;
            return ge::GRAPH_FAILED;
        }

        Col2ImCustomTilingData tiling;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());

        context->SetBlockDim(CORE_NUM);
        std::cout << "=== Col2ImCustom Tiling Start ===" << std::endl;

        // 获取输入shape: [N, C*kH*kW, L]
        const auto InputShape = context->GetInputTensor(0)->GetOriginShape();
        auto dimNum = InputShape.GetDimNum();
        
        std::cout << "Input dimension number: " << dimNum << std::endl;
        
        if (dimNum != 3) {
            std::cout << "ERROR: Input must be 3D [N, C*kH*kW, L]!" << std::endl;
            return ge::GRAPH_FAILED;
        }

        int32_t N = InputShape.GetDim(0);
        int32_t input_channels = InputShape.GetDim(1); // C*kH*kW
        int32_t L = InputShape.GetDim(2); // out_H * out_W

        std::cout << "Input shape - N: " << N 
                  << ", input_channels: " << input_channels 
                  << ", L: " << L << std::endl;

        // 获取卷积参数
        int32_t kernel_h = 2;
        int32_t kernel_w = 2;
        int32_t stride_val = 1;
        int32_t padding_val = 0;
        int32_t dilation_val = 1;
        int32_t H = 0; // 输出高度
        int32_t W = 0; // 输出宽度

        auto attrPtr = context->GetAttrs();
        if (attrPtr) {
            if (attrPtr->GetInt(0)) {
                H = static_cast<int32_t>(*(attrPtr->GetInt(0)));
            }
            if (attrPtr->GetInt(1)) {
                W = static_cast<int32_t>(*(attrPtr->GetInt(1)));
            }
            if (attrPtr->GetInt(2)) {
                kernel_h = static_cast<int32_t>(*(attrPtr->GetInt(2)));
            }
            if (attrPtr->GetInt(3)) {
                kernel_w = static_cast<int32_t>(*(attrPtr->GetInt(3)));
            }
            if (attrPtr->GetInt(4)) {
                stride_val = static_cast<int32_t>(*(attrPtr->GetInt(4)));
            }
            if (attrPtr->GetInt(5)) {
                padding_val = static_cast<int32_t>(*(attrPtr->GetInt(5)));
            }
            if (attrPtr->GetInt(6)) {
                dilation_val = static_cast<int32_t>(*(attrPtr->GetInt(6)));
            }
        }

        std::cout << "Output size - H: " << H << ", W: " << W << std::endl;
        std::cout << "Convolution params - kernel_h: " << kernel_h 
                  << ", kernel_w: " << kernel_w 
                  << ", stride: " << stride_val 
                  << ", padding: " << padding_val 
                  << ", dilation: " << dilation_val << std::endl;

        // 计算C和验证参数
        int32_t kernel_size = kernel_h * kernel_w;
        if (input_channels % kernel_size != 0) {
            std::cout << "ERROR: input_channels must be divisible by kernel_size!" << std::endl;
            return ge::GRAPH_FAILED;
        }
        int32_t C = input_channels / kernel_size;

        // 验证L的合理性
        int32_t out_H = (H + 2 * padding_val - dilation_val * (kernel_h - 1) - 1) / stride_val + 1;
        int32_t out_W = (W + 2 * padding_val - dilation_val * (kernel_w - 1) - 1) / stride_val + 1;
        int32_t expected_L = out_H * out_W;
        
        if (L != expected_L) {
            std::cout << "WARNING: L mismatch! Expected: " << expected_L << ", Got: " << L << std::endl;
        }

        std::cout << "Derived - C: " << C << ", out_H: " << out_H << ", out_W: " << out_W << std::endl;

        uint64_t ubSize = 0;
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
        std::cout << "UB Size: " << ubSize << " bytes" << std::endl;

        uint64_t input_elements = static_cast<uint64_t>(N) * input_channels * L;
        uint64_t output_elements = static_cast<uint64_t>(N) * C * H * W;

        if (output_elements == 0) {
            std::cout << "ERROR: Output elements is 0!" << std::endl;
            return ge::GRAPH_FAILED;
        }

        std::cout << "Total elements - input: " << input_elements << ", output: " << output_elements << std::endl;

        // ======= 核间划分策略 =======
        // Col2Im按输出像素划分，每个像素需要累加多个输入col值
        std::cout << "=== Inter-Core Tiling Strategy ===" << std::endl;
        
        uint32_t total_output_pixels = output_elements;
        
        int32_t base_pixels_per_core = total_output_pixels / CORE_NUM;
        int32_t big_core_num = total_output_pixels % CORE_NUM;
        int32_t big_core_pixels = base_pixels_per_core + 1;
        int32_t small_core_num = CORE_NUM - big_core_num;
        int32_t small_core_pixels = base_pixels_per_core;

        std::cout << "Base pixels per core: " << base_pixels_per_core << std::endl;
        std::cout << "Big core count: " << big_core_num << " (pixels: " << big_core_pixels << ")" << std::endl;
        std::cout << "Small core count: " << small_core_num << " (pixels: " << small_core_pixels << ")" << std::endl;

        int32_t* core_pixel_start = new int32_t[CORE_NUM];
        int32_t* core_pixel_end = new int32_t[CORE_NUM];
        int32_t* core_pixel_count = new int32_t[CORE_NUM];

        auto cleanup = [&]() {
            delete[] core_pixel_start;
            delete[] core_pixel_end;
            delete[] core_pixel_count;
        };

        int32_t current_pixel = 0;
        for (int32_t core_id = 0; core_id < CORE_NUM; core_id++) {
            if (core_id < big_core_num) {
                core_pixel_count[core_id] = big_core_pixels;
            } else {
                core_pixel_count[core_id] = small_core_pixels;
            }
            core_pixel_start[core_id] = current_pixel;
            
            if (core_pixel_count[core_id] > 0) {
                core_pixel_end[core_id] = current_pixel + core_pixel_count[core_id] - 1;
                current_pixel += core_pixel_count[core_id];
            } 
            else {
                core_pixel_end[core_id] = current_pixel - 1;
            }
        }

        std::cout << "Core distribution details:" << std::endl;
        for (int32_t core_id = 0; core_id < CORE_NUM; core_id++) {
            std::cout << "  Core " << core_id << ": start=" << core_pixel_start[core_id] 
                      << ", end=" << core_pixel_end[core_id] 
                      << ", count=" << core_pixel_count[core_id] << std::endl;
        }

        // ====== 核内划分策略 ======
        std::cout << "=== Intra-Core Tiling Strategy ===" << std::endl;
        
        int32_t tile_pixel_num = 1;

        const auto inputDataType = context->GetInputTensor(0)->GetDataType();
        uint32_t typeSize = 0;
        switch (inputDataType) {
            case ge::DT_FLOAT:
                typeSize = 4;
                break;
            case ge::DT_FLOAT16:
                typeSize = 2;
                break;
            case ge::DT_INT8:
                typeSize = 1;
                break;
            default:
                std::cout << "ERROR: Unsupported data type!" << std::endl;
                cleanup();
                return ge::GRAPH_FAILED;
        }
        std::cout << "Data type size: " << typeSize << " bytes" << std::endl;

        // UB内存需求估算
        auto EstimateUBUsage = [&](int32_t pixel_num) -> uint32_t {
            // 输出缓冲区（双缓冲）+ 累加缓冲区
            uint32_t outputBytes = AlignUp(pixel_num * typeSize, BLOCK_SIZE);
            uint32_t total = BUFFER_NUM * outputBytes + outputBytes; // 额外的累加缓冲
            
            std::cout << "  UB Estimate - pixel_num: " << pixel_num 
                      << ", output: " << outputBytes << "B"
                      << ", total: " << total << "B" << std::endl;
            
            return total;
        };

        std::cout << "UB usage optimization:" << std::endl;
        int32_t max_pixels_per_core = base_pixels_per_core + (big_core_num > 0 ? 1 : 0);
        std::cout << "Max pixels per core: " << max_pixels_per_core << std::endl;
        
        while (tile_pixel_num * 2 <= max_pixels_per_core && 
               EstimateUBUsage(tile_pixel_num * 2) <= ubSize * 85 / 100) { // 85%阈值，留更多余量
            tile_pixel_num *= 2;
            std::cout << "Increased tile_pixel_num to: " << tile_pixel_num << std::endl;
        }

        if (tile_pixel_num < MIN_TILE_SIZE) {
            tile_pixel_num = std::min(static_cast<int32_t>(MIN_TILE_SIZE), max_pixels_per_core);
            std::cout << "Adjusted tile_pixel_num to minimum: " << tile_pixel_num << std::endl;
        }

        std::cout << "Final tile_pixel_num: " << tile_pixel_num << std::endl;

        int32_t* core_loop_times = new int32_t[CORE_NUM];
        int32_t* core_tail_pixels = new int32_t[CORE_NUM];

        for (int32_t core_id = 0; core_id < CORE_NUM; core_id++) {
            if (core_pixel_count[core_id] > 0) {
                core_loop_times[core_id] = (core_pixel_count[core_id] + tile_pixel_num - 1) / tile_pixel_num;
                core_tail_pixels[core_id] = core_pixel_count[core_id] % tile_pixel_num;
                if (core_tail_pixels[core_id] == 0) {
                    core_tail_pixels[core_id] = tile_pixel_num;
                }
            } else {
                core_loop_times[core_id] = 0;
                core_tail_pixels[core_id] = 0;
            }
            std::cout << "Core " << core_id << " loop times: " << core_loop_times[core_id] 
                      << ", tail: " << core_tail_pixels[core_id] << std::endl;
        }

        auto cleanup_all = [&]() {
            cleanup();
            delete[] core_loop_times;
            delete[] core_tail_pixels;
        };

        int32_t aligned_pixel_size = ((tile_pixel_num + 32 / (int32_t)typeSize - 1) / (32 / (int32_t)typeSize)) * (32 / (int32_t)typeSize);
        std::cout << "Aligned pixel size: " << aligned_pixel_size << std::endl;
        
        // ======= 最终tiling参数汇总 =======
        std::cout << "=== Final Tiling Parameters ===" << std::endl;
        
        tiling.set_N(N);
        tiling.set_C(C);
        tiling.set_H(H);
        tiling.set_W(W);
        tiling.set_kernel_h(kernel_h);
        tiling.set_kernel_w(kernel_w);
        tiling.set_stride_val(stride_val);
        tiling.set_padding_val(padding_val);
        tiling.set_dilation_val(dilation_val);
        tiling.set_input_channels(input_channels);
        tiling.set_L(L);
        tiling.set_out_H(out_H);
        tiling.set_out_W(out_W);

        tiling.set_input_elements(input_elements);
        tiling.set_output_elements(output_elements);
        tiling.set_total_output_pixels(total_output_pixels);
        tiling.set_base_pixels_per_core(base_pixels_per_core);
        tiling.set_big_core_num(big_core_num);
        tiling.set_core_pixel_start(core_pixel_start);
        tiling.set_core_pixel_end(core_pixel_end);
        tiling.set_core_pixel_count(core_pixel_count);

        tiling.set_tile_pixel_num(tile_pixel_num);
        tiling.set_core_loop_times(core_loop_times);
        tiling.set_core_tail_pixels(core_tail_pixels);
        tiling.set_aligned_pixel_size(aligned_pixel_size);

        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        if (!currentWorkspace) { 
            std::cout << "ERROR: Failed to get workspace!" << std::endl;
            cleanup_all();
            return ge::GRAPH_FAILED;
        }
        currentWorkspace[0] = 0;
        std::cout << "Workspace size: " << currentWorkspace[0] << std::endl;

        auto tilingData = context->GetRawTilingData();
        if (!tilingData) {
            std::cout << "ERROR: Failed to get tiling data!" << std::endl;
            cleanup_all();
            return ge::GRAPH_FAILED;
        }
        
        tiling.SaveToBuffer(tilingData->GetData(), tilingData->GetCapacity());
        tilingData->SetDataSize(tiling.GetDataSize());

        std::cout << "=== Tiling Data Validation ===" << std::endl;
        std::cout << "Tiling data size: " << tiling.GetDataSize() << std::endl;
        std::cout << "=== Col2ImCustom Tiling Completed Successfully ===" << std::endl;
        
        cleanup_all();
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
static graphStatus InferShape(gert::InferShapeContext *context)
{
    if (!context) {
        return ge::GRAPH_FAILED;
    }

    const gert::Shape *InputShape = context->GetInputShape(0);
    gert::Shape *OutputShape = context->GetOutputShape(0);
    
    if (!InputShape || !OutputShape) {
        return ge::GRAPH_FAILED;
    }

    // 输入: [N, C*kH*kW, L]
    if (InputShape->GetDimNum() != 3) {
        return ge::GRAPH_FAILED;
    }

    int32_t N = InputShape->GetDim(0);
    int32_t input_channels = InputShape->GetDim(1);
    
    // 获取输出尺寸和卷积参数
    int32_t H = 0;
    int32_t W = 0;
    int32_t kernel_h = 2;
    int32_t kernel_w = 2;

    auto attrPtr = context->GetAttrs();
    if (attrPtr) {
        if (attrPtr->GetInt(0)) {
            H = static_cast<int32_t>(*(attrPtr->GetInt(0)));
        }
        if (attrPtr->GetInt(1)) {
            W = static_cast<int32_t>(*(attrPtr->GetInt(1)));
        }
        if (attrPtr->GetInt(2)) {
            kernel_h = static_cast<int32_t>(*(attrPtr->GetInt(2)));
        }
        if (attrPtr->GetInt(3)) {
            kernel_w = static_cast<int32_t>(*(attrPtr->GetInt(3)));
        }
    }

    int32_t kernel_size = kernel_h * kernel_w;
    int32_t C = input_channels / kernel_size;

    // 输出: [N, C, H, W]
    std::vector<int64_t> outputDims = {N, C, H, W};

    OutputShape->SetDimNum(outputDims.size());
    for (size_t i = 0; i < outputDims.size(); ++i) {
        OutputShape->SetDim(i, outputDims[i]);
    }

    return ge::GRAPH_SUCCESS;
}

static graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    if (!context) {
        return ge::GRAPH_FAILED;
    }

    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}

} // namespace ge

namespace ops {
class Col2ImCustom : public OpDef {
public:
    explicit Col2ImCustom(const char *name) : OpDef(name)
    {
        this->Input("col")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        
        // 定义属性：输出尺寸 + 卷积参数
        this->Attr("output_h")
            .AttrType(REQUIRED)
            .Int();
        this->Attr("output_w")
            .AttrType(REQUIRED)
            .Int();
        this->Attr("kernel_h")
            .AttrType(OPTIONAL)
            .Int(2);
        this->Attr("kernel_w")
            .AttrType(OPTIONAL)
            .Int(2);
        this->Attr("stride_val")
            .AttrType(OPTIONAL)
            .Int(1);
        this->Attr("padding_val")
            .AttrType(OPTIONAL)
            .Int(0);
        this->Attr("dilation_val")
            .AttrType(OPTIONAL)
            .Int(1);

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend910")
            .AddConfig("ascend310p")
            .AddConfig("ascend310b")
            .AddConfig("ascend910b");
    }
};
OP_ADD(Col2ImCustom);
} // namespace ops